package com.shujia.dws

import com.shujia.common.utils.DateUtils.diff_date_time
import com.shujia.common.utils.Geography.calculateLength
import com.shujia.common.utils.{DateUtils, SparkTool}
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.{DataFrame, Dataset, Row, SaveMode, SparkSession}

object DwsSpacetimeCompanionMskDay extends SparkTool {
  override def run(spark: SparkSession): Unit = {
    import spark.implicits._
    import org.apache.spark.sql.functions._

    // 从MySQL获取确诊人员名单
    val confirmedDF: DataFrame = spark
      .read
      .format("jdbc")
      .option("url", "jdbc:mysql://master:3306/crm")
      .option("dbtable", "confirmed")
      .option("user", "root")
      .option("password", "123456")
      .load()

    confirmedDF.cache()

    // 加载位置数据融合表
    val mergeLocDF: DataFrame = spark.table("dwi.dwi_res_regn_mergelocation_msk_d").where($"day_id" === day_id)

    // 从确诊人员名单中提取手机号，并转成本地集合
    val confirmMdnArr: Array[String] = confirmedDF.select($"mdn").as[String].collect()

    // 广播集合
    val confirmMdnArrBro: Broadcast[Array[String]] = spark.sparkContext.broadcast(confirmMdnArr)

    // 从位置数据融合表中筛选出确诊人员的行动轨迹
    val confirmLoc: DataFrame = mergeLocDF
      .filter(row => confirmMdnArrBro.value.contains(row.getAs[String]("mdn")))
      .select($"mdn" as "c_mdn", $"start_date" as "c_start_date", $"longi" as "c_longi", $"lati" as "c_lati", $"county_id")

    /**
     * 时空伴随者：
     * 1、时间不超过1个小时
     * 2、距离不超过500m
     */

    /**
     * 优化：（尽可能地减少比较次数）
     * 1、正常人员的行动轨迹同确诊人员的行动轨迹比较时，只要有一条纪录符合时空伴随者的条件，那后续的确诊人员行动轨迹的判断就可以跳过
     * 2、将正常人员一天的行动轨迹按照手机号mdn进行分组，再与确诊人员的行动轨迹进行对比，只要有一条纪录符合时空伴随者的条件，就可以直接跳过当前人员，取下一个人员的记录继续比较
     * 3、将属于同一个区县的行动轨迹进行对比，筛选出时空伴随者
     */

    val cMdnDF: DataFrame = confirmedDF.select($"mdn" as "confirm_mdn")

    val spaceTimeDF: DataFrame = mergeLocDF
      // 过滤掉确诊人员的行动轨迹，保留正常人的行动轨迹
      .join(cMdnDF, $"mdn" === $"confirm_mdn", "left")
      .where($"confirm_mdn".isNull) // 当未关联上时需要注意右表的数据为NULL，不能直接判断相等/不相等
      // 使用county_id区县id进行关联
      .join(confirmLoc, "county_id")
      // 逐条数据进行判断是否为时空伴随者
      // 计算时间差
      .withColumn("diff_time", diffTime($"start_date", $"c_start_date"))
      .where($"diff_time" < 3600)
      // 计算距离
      .withColumn("distance", calculateLength($"longi", $"lati", $"c_longi", $"c_lati"))
      .where($"distance" < 500)
      .select(
        $"mdn"
        , $"start_date"
        , $"end_date"
        , $"county_id"
        , $"longi"
        , $"lati"
        , $"bsid"
        , $"grid_id"
        , $"c_mdn"
        , $"c_start_date"
        , $"c_longi"
        , $"c_lati"
      )

    spaceTimeDF.cache()

    spaceTimeDF
      .write
      .format("csv")
      .option("sep", ",")
      .mode(SaveMode.Overwrite)
      .save(s"/daas/motl/dws/dws_spacetime_companion_msk_d/day_id=$day_id")

    // 添加分区
    spark.sql(
      s"""
         |alter table dws.dws_spacetime_companion_msk_d  add if not exists partition(day_id='$day_id')
         |""".stripMargin)

    // 提取时空伴随者的手机号 进行去重
    spaceTimeDF
      .select($"mdn")
      .distinct()
      .write
      .format("csv")
      .option("sep", ",")
      .mode(SaveMode.Overwrite)
      .save(s"/daas/motl/dws/dws_spacetime_companion_mdn_msk_d/day_id=$day_id")


    // 添加分区
    spark.sql(
      s"""
         |alter table dws.dws_spacetime_companion_mdn_msk_d  add if not exists partition(day_id='$day_id')
         |""".stripMargin)

  }


}
